from django.shortcuts import render
from django.core.files.storage import FileSystemStorage
from django.conf import settings
from predict import load_model, predict_image
import os

model = load_model()

def load_class_names(file_path="name_list.txt"):
    class_names = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split(" ", 1)
            if len(parts) == 2:
                class_id, artist_name = parts
                class_names.append(artist_name)
    return class_names

class_names = load_class_names(os.path.join(settings.BASE_DIR, 'name_list.txt'))


def load_author_info(file_path="author_info.txt"):
    author_info = {}
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
        i = 0
        current_author = None
        current_info = []

        while i < len(lines):
            line = lines[i].strip()

            if line:
                if current_author is None:
                    if current_info:
                        author_info[current_author] = "\n".join(current_info)

                    current_author = line
                    current_info = []
                else:
                    current_info.append(line)
            else:
                if current_author and current_info:
                    author_info[current_author] = "\n".join(current_info)
                    current_author = None
                    current_info = []
            i += 1

        if current_author and current_info:
            author_info[current_author] = "\n".join(current_info)

    return author_info

author_info = load_author_info(os.path.join(settings.BASE_DIR, 'author_info.txt'))


def upload_artwork(request):
    clear_media_folder()
    if request.method == 'POST' and request.FILES.get('image'):
        uploaded_file = request.FILES['image']
        fs = FileSystemStorage()
        file_path = os.path.join(settings.MEDIA_ROOT, uploaded_file.name)
        fs.save(file_path, uploaded_file)
        with open(file_path, 'rb') as f:
            image_bytes = f.read()
        predicted_class = predict_image(model, image_bytes)
        predicted_artist = class_names[predicted_class]

        artist_info = author_info.get(predicted_artist, "未找到该作者的详细信息")
        if artist_info.startswith(predicted_artist + ":"):
            artist_info = artist_info[len(predicted_artist) + 1:].strip()

        return render(request, 'success.html', {
            'author': predicted_artist,
            'author_info': artist_info,
            'image_url': fs.url(uploaded_file.name),
            'class_names': class_names
        })
    return render(request, 'index.html', {'class_names': class_names})

def clear_media_folder():
    media_folder = settings.MEDIA_ROOT
    for filename in os.listdir(media_folder):
        file_path = os.path.join(media_folder, filename)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)
        except Exception as e:
            print(f"Error deleting file {file_path}: {e}")